from config import config_singleton
import time
import socket
import struct
from utils import sys_logger
import queue
import threading
import keyboard
from utils import BaseClient


class Udp_Multicast_Receiver(BaseClient):

    """
    This class is designed to receive UDP multicast messages and manage them using a queue.
    It also provides methods to control the reception and handling of these messages.
    multicast_group: str,
    port: int,
    queue_size: int
    """

    def __init__(self, multicast_group, port, queue_size, isSend = False):
       
        self.broadcast_addr = (multicast_group, port)
        # 创建一个socket
        self._sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM, socket.IPPROTO_UDP)
        # 允许端口复用
        self._sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        # 设置接收/发送缓存为1M
        self._sock.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 2048*2048)
        self._sock.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 1024*1024)
        
        if not isSend:
            self._sock.bind(("", port))

        # 加入组播组
        mreq = struct.pack("4sl", socket.inet_aton(multicast_group), socket.INADDR_ANY)
        self._sock.setsockopt(socket.IPPROTO_IP, socket.IP_ADD_MEMBERSHIP, mreq)
        # 创建控制事件
        # self._sock.setblocking(False)

        self._recv_queue = queue.Queue(maxsize=queue_size)
        self._send_queue = queue.Queue(maxsize=queue_size)

        self.need_stop = False
        self.byte_cnt = 0
        self.recv_buffer_size = 128 * 1024

    def do_receive(self):

        while not self.need_stop:
            try:
                data, addr = self._sock.recvfrom(self.recv_buffer_size)

                # if data == b'stop':
                #     sys_logger.info("Udp Receiver stopped expectedly.")
                #     break

                if self._recv_queue.full():
                    sys_logger.debug(f"{len(data)} bytes data droped, udp recv queue is too small, set queue maxsize more bigger.")
                else:
                    self._recv_queue.put(data)
                
                self.byte_cnt += len(data)

            except BlockingIOError as e:
                # 没有数据可读，等待一段时间后再重试
                time.sleep(0.1)
            except OSError as e:
                if self.need_stop:
                    sys_logger.info(f"Udp Receiver stopped expectedly. Total recv bytes {self.byte_cnt}")
                else:
                    sys_logger.error(f"Udp Receiver stopped unexpectedly due to an error. Total recv bytes {self.byte_cnt}")
                    sys_logger.error(f"recv errs:{e.strerror}")
                    break
            finally:
                pass

    def do_send(self):
        """
        """
        while not self.need_stop:
            try:
                data = self._send_queue.get()

                if data is None:
                    break
                
                # 发送至广播地址
                self._sock.sendto(data, self.broadcast_addr)

            except socket.error as e:
                sys_logger.error(f"Failed to send data to udp client : {e}")

    def get_recv_data(self):
        # 如果没有数据，自动进入阻塞模式
        return self._recv_queue.get()
    
    def put_send_data(self, data):
        """ 直接将data放入发送队列 """
        self._send_queue.put(data)  


    def stop(self):
        self.need_stop = True
        # self._sock.sendto(b'stop', self._sock.getsockname())       # 向自己发送数据，激活sock.recv()
        self._sock.shutdown(socket.SHUT_RDWR)
        self._sock.close()
        self._recv_queue.put(None)  # 停止receiving
        self._send_queue.put(None)  # 停止sending


# test example
if __name__ == "__main__":

    # 测试可行性
    # read configs
    ip = config_singleton.udp_setting["ip"]
    port = config_singleton.udp_setting["port"]
    maxsize = config_singleton.udp_setting["queue_maxsize"]
    udp_recv = Udp_Multicast_Receiver(ip, port, maxsize)

    thr = threading.Thread(target=udp_recv.do_receive)
    thr.start()
    keyboard.wait('esc')
    udp_recv.stop()
